# -*- coding: utf-8 -*-
import os
import math
import re
import time
import jieba
from config import dict_file, stop_word_file, sentence_split_pattern
from utils import get_stopwords, load_dictionary, cut_sentence, generate_ngram

class Node(object):
    """
    建立字典树的节点
    """

    def __init__(self, word=None, frequence=0, parent=None):
        # 记录词汇短语字符串
        self.word = word

        # 父节点
        self.parent = parent

        # 后子节点,采用dict数据类型，初始化为空
        self.suffix = {}

        # 仅限用于第二级子节点，例如AB组合词，在计算左熵时，出现C->AB之类词汇，为了记录方便，所以用此变量记录左词节点C，相当于C->A->B,采用dict数据类型，初始化为空
        self.prefix = {}

        # 词频计数
        self.frequence = frequence

        # 右子词频总计数
        self.suffix_frequence = 0

        # 仅限用于第二级子节点，左子词频总计数
        self.prefix_frequence = 0

        # PMI  互信息值
        self.PMI = 0

        # PL 左熵
        self.PL = 0

        # PR 右熵
        self.PR = 0

        # SC 综合分
        self.SC = 0


class PhraseDiscovery(object):
    """
    新词发现工具，对node数据进行分析管理，并且包含统计词频，计算左右熵，计算互信息的方法
    """

    def __init__(self):
        """
        初始discovery
        """

        # 建立根节点
        self.root = Node()

        # 加载分词字符
        self.stopwords = get_stopwords(stop_word_file)

        # 加载通用词频库
        word_freq = load_dictionary(dict_file)
        for word, frequence in word_freq.items():
            new_node = Node(word, int(frequence), self.root)
            self.root.suffix[word] = new_node

            # 汇总到上一级, 注意:上一级的frequence不用叠加，根据ngram词组组合的计算已经单独添加完毕,只在初始化词汇时用
            self.root.suffix_frequence += new_node.frequence

    def load_file_split_words(self, filename):
        """
        :param filename:资料文件
        :return: 二维数组,[[句子1分词list], [句子2分词list],...,[句子n分词list]]
        """
        start = time.time()
        print('------> get_line_split_words start...')
        data = []
        with open(filename, 'r', encoding='utf-8') as f:
            for line in f:
                word_list = [x for x in jieba.cut(line.strip(), cut_all=False) if x not in self.stopwords]
                data.append(word_list)
        f.close()
        print('------> get_line_split_words end...', 'time last:', time.time() - start)

        return data

    def add_line_ngram(self, data):
        """
        生成1到3阶词组合列表ngram，添加到discovery计算模型
        :param data:
        :return:
        """
        start = time.time()
        print('------> add_line_ngram start...')
        for word_list in data:
            # tmp 表示每一行自由组合后的结果（n gram）
            # tmp: [['它'], ['是'], ['小'], ['狗'], ['它', '是'], ['是', '小'], ['小', '狗'], ['它', '是', '小'], ['是', '小', '狗']]
            ngrams = generate_ngram(word_list, 3)
            for d in ngrams:
                self.add_words(d)
        print('------> add_line_ngram end...', 'time last:', time.time() - start)

    def add_words(self, words):
        """
        基础方法：添加词汇组合节点，并计数
        :param words:
        :return:  相当于对 [a, b, c] a->b->c, [b, c, a] b->c->a
        """

        length = len(words)
        node = self.root

        # suffix 添加节点并计数
        for index, word in enumerate(words):
            # found_in_child = False
            # 根据字符key找子节点
            subnode = node.suffix.get(word)

            if subnode is None:
                new_node = Node(word, parent=node)
                # new_node.parent = node
                node.suffix[word] = new_node
                node = new_node  # 将新子节点转换为当前节点
            else:
                node = subnode  # 将找到子节点转换为当前节点

            # 判断是否是最后一个节点，这个词每出现一次就+1
            if index == length - 1:
                node.frequence += 1
                node.parent.suffix_frequence += 1  # 汇总到上一级,无需向上递归
                if length == 2:  # 如果是二阶，则将二阶词汇计数统计到根节点的 self.root.frequence （这是个trick,根节点frequence无其他用处）
                    self.root.frequence += 1

        # prefix 添加节点并计数，只作用于三级词汇，比如abc,a是bc的左熵节点,为了便于计算记录到b->c的子节点上，记录为b->c->a
        node = self.root
        if length == 3:
            words = list(words)
            words[0], words[1], words[2] = words[1], words[2], words[0]

            for index, word in enumerate(words):
                if index == 2:  # 第三级在prefix中找节点
                    subnode = node.prefix.get(word)
                else:  # 不是第三级在suffix中找节点
                    subnode = node.suffix.get(word)

                if subnode is None:
                    new_node = Node(word)
                    new_node.parent = node
                    node.prefix[word] = new_node
                    node = new_node  # 将新子节点转换为当前节点
                else:
                    node = subnode  # 将找到子节点转换为当前节点


                # 判断是否是最后一个节点，这个词每出现一次就+1
                if index == 2:
                    node.frequence += 1
                    node.parent.prefix_frequence += 1  # 汇总到上一级,无需向上递归


    def analyse(self, PMI_limit=0):
        """
        基础方法
        添加完词汇后，最后进行综合计算分析
        可以分多步加载数据，最后一步分析，不影响结果
        遍历计算所有新出现的二级词汇信息熵，
        # 计算公式 score = PMI + min(左熵， 右熵) => 熵越小，说明越有序，这词再一次可能性更大！
        :return: 返回的分析结果变量
        """
        start = time.time()
        print('------>analyse start.......')

        # 需要返回的分析结果变量
        newphrases = []

        # 一阶词总计数变量
        total = self.root.suffix_frequence

        # 二阶词总计数变量
        total2 = self.root.frequence

        # 遍历计算所有新出现的二级词汇信息熵，
        for node1 in self.root.suffix.values():
            for node2 in node1.suffix.values():
                # 先计算二阶词的互熵,  互信息值越大，说明 a,b 两个词相关性越大
                char1_ratio = node1.frequence / total  # 第一词汇概率
                char2_ratio = self.root.suffix[node2.word].frequence / total  # 第二词汇概率
                node2.PMI = math.log(max(node2.frequence, 1), 2) - math.log(total2, 2) - \
                            math.log(char1_ratio, 2) - math.log(char2_ratio, 2)

                # 计算左熵
                for node3 in node2.prefix.values():
                    node2.PL -= (node3.frequence / node2.prefix_frequence) * math.log(
                        node3.frequence / node2.prefix_frequence, 2)

                # 计算右熵
                for node3 in node2.suffix.values():
                    node2.PR -= (node3.frequence / node2.suffix_frequence) * math.log(
                        node3.frequence / node2.suffix_frequence, 2)

                # 计算综合熵 ,计算公式 score = PMI + min(左熵， 右熵) => 熵越小，说明越有序，这词再一次可能性更大！
                node2_ratio = node2.frequence / total2  # node2出现概率
                node2.SC = (node2.PMI + min(node2.PL, node2.PR)) * node2_ratio

                # 这里做了PMI阈值约束
                if node2.PMI > PMI_limit:
                    newphrases.append(node2)

        print('------>analyse end.......', 'time last:', time.time() - start)
        return newphrases

    def add_parse_file0(self, textfile):
        """
        添加文本文件数据，每行一句话sentence
        :return:None
        :param filename:资料文件
        :return: 二维数组,[[句子1分词list], [句子2分词list],...,[句子n分词list]]
        """
        start = time.time()
        print('------> add_parse_file0 start...')
        data = []
        with open(textfile, 'r', encoding='utf-8') as f:
            for line in f:
                # 第一步：分词得到 word_list=['它', '是', '小', '狗']
                word_list = [x for x in jieba.cut(line.strip(), cut_all=False) if x not in self.stopwords]

                # 第二步： ngrams 表示每一行自由组合后的结果（n gram）
                # ngrams: [['它'], ['是'], ['小'], ['狗'], ['它', '是'], ['是', '小'], ['小', '狗'], ['它', '是', '小'], ['是', '小', '狗']]
                ngrams = generate_ngram(word_list, 3)

                # 第三步：加入到node数据模型中
                for d in ngrams:
                    self.add_words(d)
        f.close()
        print('------> add_parse_file0 end...', 'time last:', time.time() - start)

    def add_parse_file(self, textfile):
        """
        添加文本文件数据
        :return:None
        :param filename:资料文件
        :return: 二维数组,[[句子1分词list], [句子2分词list],...,[句子n分词list]]
        """
        start = time.time()
        print('------> add_parse_file start...')

        with open(textfile, 'r', encoding='utf-8') as f:
            for line in f:
                # 第一步：将line，也就是段落paragraph，分成句子，得到 sentence_list=['它是小狗，','不是狐狸。']
                sentences = cut_sentence(line.strip(), sentence_split_pattern)

                for sentence in sentences:
                    # 第二步：分词得到 word_list=['它', '是', '小', '狗']
                    word_list = [x for x in jieba.cut(sentence.strip(), cut_all=False) if x not in self.stopwords]

                    # 第三步： ngrams 表示每一行自由组合后的结果（n gram）
                    # ngrams: [['它'], ['是'], ['小'], ['狗'], ['它', '是'], ['是', '小'], ['小', '狗'], ['它', '是', '小'], ['是', '小', '狗']]
                    ngrams = generate_ngram(word_list, 3)

                    # 第四步：加入到node数据模型中
                    for d in ngrams:
                        self.add_words(d)
        f.close()
        print('------> add_parse_file end...', 'time last:', time.time() - start)

    def add_parse_text(self, text):
        """
        直接添加文本数据
        :return:None
        """
        start = time.time()
        print('------> add_parse_text start...')

        # 第一步：将line，也就是段落paragraph，分成句子，得到 sentence_list=['它是小狗。','不是狐狸。']
        sentences = cut_sentence(text.strip())

        for sentence in enumerate(sentences):
            # 第二步：分词得到 word_list=['它', '是', '小', '狗']
            word_list = [x for x in jieba.cut(sentence.strip(), cut_all=False) if x not in self.stopwords]

            # 第三步： ngrams 表示每一行自由组合后的结果（n gram）
            # ngrams: [['它'], ['是'], ['小'], ['狗'], ['它', '是'], ['是', '小'], ['小', '狗'], ['它', '是', '小'], ['是', '小', '狗']]
            ngrams = generate_ngram(word_list, 3)

            # 第四步：加入到node数据模型中
            for d in ngrams:
                self.add_words(d)

        print('------> add_parse_text end...', 'time last:', time.time() - start)


    def discovery_phrase(self, topnum=100, PMI_limit=0):
        # 执行分析方法
        newphrases = self.analyse(PMI_limit)

        start = time.time()
        print('------>words sort start.......')

        # 按照 大到小倒序排列，value 值越大，说明是组合词的概率越大
        # result变成 => [('世界卫生_大会', 0.4380419441616299), ('蔡_英文', 0.28882968751888893) ..]
        result = sorted(newphrases, key=lambda x: x.SC, reverse=True)

        print('------>words sort end.......', 'time last:', time.time() - start)

        return result[:topnum]

